import sqlite3
from Parameters import *
from PolymorphismTables import PolymorphismTables
from R2_Tools import R2
from SiteFrequencySpectrum import SiteFrequencySpectrum
import infiniteSitesTools as ist 
from Polyspect_PowerSphere import PolyspectPowerSphere
from Databases import DB_Polyspect

import math
#import matplotlib.pyplot as plt
import numpy.matlib as np

import sys

def main(argv):
	
	print "***************************************************"
	print "********************* New run *********************"
	print "***************************************************"
	print ""
	
	nChr = int(argv[0])
	segSites = int(argv[1])
	alpha = float(argv[2])

	ptNull = None
	mySFS = None

	alphaToCheck = [ alpha ] #, 2.0, 5.0, 10.0, 20.0, 50.0, 100.0]
	tauToCheck = [ 1000, 500, 200, 100, 10000, 5000, 2000 ]

	for alpha in alphaToCheck:
		 for tau in tauToCheck:

			pop = getDefaultPopulationParameters( nChr, segSites )
			
			N1 = pop.N * math.exp(-alpha * tau / (4.0*pop.N))
			print "N1\t", N1
			
			nullDemography = Demography(pop.N)
			nullDemography.addEpoch( alpha, 0, False)
			nullDemography.addEpoch( N1, tau, True )


			print "Ho\t", nullDemography
			print pop.getMSString_Null(), nullDemography.getMSString()

			simParams = SimulationParameters( pop, nullDemography )
			print simParams
						
			altDemography = Demography( pop.N )
			altDemography.addEpoch( pop.N, 0, True )
			print "Ha\t", altDemography
		
			ptNull = PolymorphismTables()
			msCommand = simParams.getMSString_Null()
			msOut = ptNull.runMS(msCommand)
			ptNull.readMS(None, msOut)
	
			mySFSs = SiteFrequencySpectrum(ptNull)
			
			pt = PolymorphismTables()
			msCommand = simParams.getMSString_Alternative(altDemography )
			msOut = pt.runMS(msCommand)
			pt.readMS(None, msOut)
			myAltSFSs = SiteFrequencySpectrum(pt)
			
			db = DB_Polyspect()
			polyspects = db.getAllPolyspects( nChr )
			
			mySpace = PolyspectPowerSphere( simParams, altDemography, mySFSs, myAltSFSs, polyspects)
			
			point = mySpace.getCurrentPoint()
			print point
			
			for i in range(0,10000):
				newPoint = mySpace.getNextPoint()
				if newPoint == None:
					print "\nFound peak for now"
					break
					
				print newPoint
		
			db.addPowerPoint( mySpace )
					
#	rr = R2( simParams )
#	r2power = rr.getPower(altDemography)
#	print 'R2 power:\t' + str(r2power)
#
#	makePlot( mySpace, r2power )


def getDefaultPopulationParameters( nChr=28, segSites=50 ):
	
	popParameters = Population()
	
	popParameters.numChromosomes = nChr
	popParameters.numSegregatingSites = segSites

	popParameters.rateMutation = 1E-8
	popParameters.rateRecombination = 0.0
	popParameters.sequenceLength = 25000

	popParameters.N = 10000
	popParameters.numSimulations = 20000
	popParameters.numPowerSimulations = 12500

	return popParameters


def makePlot( mySpace, r2power ):
	tjd = ist.normalizeToFirst( mySpace.myPowerPath[0].myPolyspect.myOperator )
	flf = ist.normalizeToFirst( mySpace.myPowerPath[3].myPolyspect.myOperator )
	dmd = ist.normalizeToFirst( mySpace.myPowerPath[-1].myPolyspect.myOperator )


	myPowers = []
	myPowers.append( mySpace.myPowerPath[0].myPower)
	myPowers.append( mySpace.myPowerPath[2].myPower)
	myPowers.append( mySpace.myPowerPath[3].myPower)
	myPowers.append( r2power )
	myPowers.append( mySpace.myPowerPath[-1].myPower)
	
	fig = plt.figure()
	ind = np.arange(30)
	
	ax = fig.add_subplot(2,2,1)
	ax.bar( ind, tjd, color='m' )
	ax.set_title('Tajimas D')

	ax = fig.add_subplot(2,2,2)
	ax.bar( ind, flf, color='b')
	ax.set_title('Fu Lis F')

	ax = fig.add_subplot(2,2,3)
	ax.bar( ind, dmd, color='y')
	ax.set_title('Davids M')

	ind = np.arange(5)
	ax = fig.add_subplot(2,2,4)
	rects = ax.bar( ind, myPowers, color='r')
	ax.set_title('Powers')
	ax.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
	ax.set_xticks(ind+0.35)
	ax.set_xticklabels( ('TajD', 'FuLiD', 'FuLiF', 'R2', 'DavidsM') )

	def autolabel(rects):
		# attach some text labels
		for rect in rects:
			height = rect.get_height()
			ax.text(rect.get_x()+rect.get_width()/2., 1.05*height, '%d'%int(100*height), ha='center', va='bottom')

	autolabel(rects)

	print "\n*** Final path ***"
	print mySpace
	print "\nProgram complete"
	
	plt.show()


	#r = codebox( "*** Final path ***", "Results", str(mySpace) )
	
if __name__ == "__main__":
	main(sys.argv[1:])
